#pragma once 
#include <set>
#include "requestor.hpp"

/*  
    客户端的功能比较分离，注册端和发现端不在同一主机上。
    1、作为服务提供者，需要一个能够进行服务注册的接口，连接注册中心，进行服务注册
    2、作为服务发现者，需要一个能够进行服务发现的接口，将获取到的能够提供指定服务的主机信息管理起来
*/

namespace yjz_rpc
{
    namespace client
    {
        class Provider
        {
        public:
            using ptr = std::shared_ptr<Provider>;
            Provider(const Requestor::ptr& requestor) : _requestor(requestor) {}

            bool serviceRegistry(const BaseConnection::ptr& conn, const std::string& method, const Address& host)
            {
                auto req_msg = MessageFactory::create<ServiceRequest>();
                req_msg->setHost(host);
                req_msg->setId(UUID::uuid());
                req_msg->setMethod(method);
                req_msg->setMType(MType::REQ_SERVICE);
                req_msg->setOpType(ServiceOpType::SERVICE_REGISTRY);

                BaseMessage::ptr rsp_msg;
                bool res = _requestor->send(conn, req_msg, rsp_msg);
                if (res == false)
                {
                    LOG(LogLevel::ERROR) << "服务注册失败： " << method;
                    return false;
                }
                // 基类无法调用返回码，必须转换为子类
                auto service_rsp = std::dynamic_pointer_cast<ServiceResponse>(rsp_msg);
                if (service_rsp.get() == nullptr)
                {
                    LOG(LogLevel::ERROR) << "响应类型向下转换失败！";
                    return false;
                }
                if (service_rsp->rcode() != RCode::RCODE_OK)
                {
                    LOG(LogLevel::ERROR) << "服务注册失败，失败原因： " << errReason(service_rsp->rcode());
                    return false;
                }
                return true;
            }
        private:
            Requestor::ptr _requestor;
        };

        // 负载均衡
        class MethodHost
        {
        public:
            using ptr = std::shared_ptr<MethodHost>;

            MethodHost() : _idx(0) {}
            MethodHost(const std::vector<Address>& host)
                : _hosts(host), _idx(0)
            {}

            // 中途服务上线
            void appendHost(const Address& host)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                _hosts.push_back(host);
            }
            
            // 中途服务下线
            void removeHost(const Address& host)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                for (auto it = _hosts.begin(); it != _hosts.end(); it++)
                {
                    if (*it == host)
                    {
                        _hosts.erase(it);
                    }
                }
            }

            Address chooseHost()
            {
                std::unique_lock<std::mutex> lock(_mutex);
                size_t pos = _idx++ % _hosts.size();
                return _hosts[pos];
            }

            bool empty() 
            {
                std::unique_lock<std::mutex> lock(_mutex);
                return _hosts.empty();
            }

        private:
            size_t _idx;
            std::mutex _mutex;
            std::vector<Address> _hosts;
        };

        class Discoverer
        {
        public:
            using ptr = std::shared_ptr<Discoverer>;
            using offlineCallback = std::function<void(const Address&)>;
            Discoverer(const Requestor::ptr& req, const offlineCallback& cb) 
                : _requestor(req) 
                , _offline_callback(cb)
            {}

            bool serviceDiscovery(const BaseConnection::ptr& conn, const std::string& method, Address& host)
            {
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it = _method_hosts.find(method);
                    if (it != _method_hosts.end())
                    {
                        if (it->second->empty() == false)
                        {
                            host = it->second->chooseHost();
                            return true;
                        }
                    }
                }
                // 当前还没有提供者提供该服务
                auto req_msg = MessageFactory::create<ServiceRequest>();
                req_msg->setId(UUID::uuid());
                req_msg->setMethod(method);
                req_msg->setMType(MType::REQ_SERVICE);
                req_msg->setOpType(ServiceOpType::SERVICE_DISCOVERY);
                BaseMessage::ptr rsp_msg;
                bool res = _requestor->send(conn, req_msg, rsp_msg);
                if (res == false)
                {
                    LOG(LogLevel::ERROR) << "服务发现失败！";
                    return false;
                }
                auto service_rsp = std::dynamic_pointer_cast<ServiceResponse>(rsp_msg);
                if (service_rsp.get() == nullptr)
                {
                    LOG(LogLevel::ERROR) << "响应类型向下转换失败！";
                    return false;
                }
                if (service_rsp->rcode() != RCode::RCODE_OK)
                {
                    LOG(LogLevel::ERROR) << "服务发现失败！ " << errReason(service_rsp->rcode());
                    return false;
                }
                std::unique_lock<std::mutex> lock(_mutex);
                auto method_host = std::make_shared<MethodHost>(service_rsp->hosts());
                if (method_host->empty())
                {
                    LOG(LogLevel::ERROR) << "服务发现失败，没有提供该服务的主机！" << method;
                    return false;
                }
                host = method_host->chooseHost();
                _method_hosts[method] = method_host;
                return true;
            }

            // 提供给Dispatcher模块进行服务上线、下线、请求处理的回调函数
            void onServiceRequest(const BaseConnection::ptr& conn, const ServiceRequest::ptr& ser_req)
            {
                auto optype = ser_req->opType();
                std::string method = ser_req->method();
                std::unique_lock<std::mutex> lock(_mutex);
                // 上线请求
                if (optype == ServiceOpType::SERVICE_ONLINE)
                {
                    auto it = _method_hosts.find(method);
                    if (it == _method_hosts.end())
                    {
                        auto method_host = std::make_shared<MethodHost>();
                        method_host->appendHost(ser_req->host());
                        _method_hosts[method] = method_host;
                    }
                    else
                    {
                        it->second->appendHost(ser_req->host());
                    }
                }
                else if (optype == ServiceOpType::SERVICE_OFFLINE)
                {
                    auto it = _method_hosts.find(method);
                    if (it == _method_hosts.end())
                    {
                        return;
                    }
                    it->second->removeHost(ser_req->host());
                    _offline_callback(ser_req->host());
                }
            }
        private:
            std::mutex _mutex;
            Requestor::ptr _requestor;
            offlineCallback _offline_callback;
            std::unordered_map<std::string, MethodHost::ptr> _method_hosts; // 该服务可以由哪些主机提供
        };
    }
}